package com.yanggu.flink.state_process_api.read

import cn.hutool.core.date.DateUtil
import org.apache.flink.api.common.restartstrategy.RestartStrategies
import org.apache.flink.api.common.state.{ListState, ListStateDescriptor, MapStateDescriptor, ValueStateDescriptor}
import org.apache.flink.runtime.state.{FunctionInitializationContext, FunctionSnapshotContext}
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction
import org.apache.flink.streaming.api.environment.CheckpointConfig.ExternalizedCheckpointCleanup
import org.apache.flink.streaming.api.functions.co.KeyedBroadcastProcessFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.streaming.api.scala.function.ProcessWindowFunction
import org.apache.flink.streaming.api.windowing.windows.GlobalWindow
import org.apache.flink.util.Collector

import java.util.Date
import scala.collection.mutable.ListBuffer

/**
 * <p>生成operator state中的list state和union state
 * <p>生成broadcast state
 * <p>生成keyed state中的list state
 * <p>生成count window
 */
object CheckpointCreate {

  def main(args: Array[String]): Unit = {

    val env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI()
    env.setParallelism(1)
    env.enableCheckpointing(60000)
    //最多支持20个检查点
    env.getCheckpointConfig.setMaxConcurrentCheckpoints(20)
    env.getCheckpointConfig.setCheckpointStorage("file:///E:/tmp/flink/checkpoint")
    //程序结束时不删除检查点
    env.getCheckpointConfig.setExternalizedCheckpointCleanup(ExternalizedCheckpointCleanup.RETAIN_ON_CANCELLATION)
    env.setRestartStrategy(RestartStrategies.fixedDelayRestart(20, 5000L))

    val socketDataStream = env.socketTextStream("localhost", 9000)

    //生成广播流
    val broadcastDataStream = socketDataStream
      .broadcast(new MapStateDescriptor[String, String]("test-broadcast-state", classOf[String], classOf[String]))

    val keyedStream = socketDataStream.keyBy(_ => "1")
    keyedStream
      .connect(broadcastDataStream)
      .process(new MyFunction).uid("test-save-read")
      .print("test-list-state")

    //生成count window state
    keyedStream
      .countWindow(5L)
      .process(new ProcessWindowFunction[String, String, String, GlobalWindow] {
        override def process(key: String, context: Context, elements: Iterable[String], out: Collector[String]): Unit = {
          out.collect(s"key: $key, maxTimeStamp: ${DateUtil.formatDateTime(new Date(context.window.maxTimestamp()))}, data: ${elements.mkString(",")}")
        }
      }).uid("test-save-read-window")
      .print("test-count-window-state")

    env.execute()

  }

}

class MyFunction extends KeyedBroadcastProcessFunction[String, String, String, Int] with CheckpointedFunction {

  private var listState: ListState[Int] = _

  private var unionListState: ListState[Int] = _

  private val listBuffer = ListBuffer[Int]()

  private val unionListBuffer = ListBuffer[Int]()

  private lazy val valueState = getRuntimeContext.getState(new ValueStateDescriptor[Int]("keyed-state-test", classOf[Int], 0))

  override def snapshotState(context: FunctionSnapshotContext): Unit = {
    //清空状态中的数据
    listState.clear()
    unionListState.clear()
    //把本地变量中的数据保存到listState中
    listBuffer.foreach(listState.add(_))
    unionListBuffer.foreach(unionListState.add(_))
  }

  override def initializeState(context: FunctionInitializationContext): Unit = {
    listState = context.getOperatorStateStore.getListState(new ListStateDescriptor[Int]("test-list-state", classOf[Int]))
    unionListState = context.getOperatorStateStore.getUnionListState(new ListStateDescriptor[Int]("test-union-list-state", classOf[Int]))
    //如果是故障恢复, 需要将checkpoint中的数据拷贝到本地缓存中
    if (context.isRestored) {
      listState.get().forEach(listBuffer.append(_))
      unionListState.get().forEach(unionListBuffer.append(_))
    }
  }

  override def processElement(value: String, ctx: KeyedBroadcastProcessFunction[String, String, String, Int]#ReadOnlyContext, out: Collector[Int]): Unit = {
    //添加到本地缓存中
    listBuffer.append(value.toInt)
    unionListBuffer.append(value.toInt)
    valueState.update(valueState.value() + value.toInt)
    out.collect(value.toInt)
  }

  override def processBroadcastElement(value: String, ctx: KeyedBroadcastProcessFunction[String, String, String, Int]#Context, out: Collector[Int]): Unit = {
    val broadcastState = ctx.getBroadcastState(new MapStateDescriptor[String, String]("test-broadcast-state", classOf[String], classOf[String]))
    broadcastState.put(value, value)
  }

}
